#include "MDictBase.h"

#include <zlib.h>

#include <cassert>
#include <regex>
#include <stdexcept>

static void printHex(const vector<uint8_t>& v) {
    for (auto i : v) {
        printf("%02hhx ", i);
    }
    cout << endl;
}
// 构造函数
MDictBase::MDictBase(
    const std::string& fname, const std::string& encoding,
    const std::pair<std::vector<uint8_t>, std::string>& passcode) {
    m_filename = fname;
    m_encoding = encoding;
    m_num_entries = 0;
    m_encrypt = 0;
    m_key_block_offset = 0;
    m_record_block_offset = 0;
    m_key_index_offset = 0;
    m_key_data_offset = 0;
    m_record_index_offset = 0;

    // 如果提供了密码，设置加密密钥
    if (!passcode.first.empty()) {
        m_encrypted_key = new uint8_t[passcode.first.size()];
        std::memcpy(m_encrypted_key, passcode.first.data(),
                    passcode.first.size());
    }

    // 读取文件头
    _read_header();

    // 读取key
    _read_keys();
    return;
    m_key_list = _read_keys();

    for (int i = 0; i < m_key_list.size(); i++) {
        uint64_t offset = m_key_list[i].first;
        uint32_t size = m_key_list[i + 1].first - offset;
        m_key_offset[m_key_list[i].second] = {offset, size};
    }
}

// 析构函数
MDictBase::~MDictBase() {
    if (m_encrypted_key) {
        delete[] m_encrypted_key;
        m_encrypted_key = nullptr;
    }
}

// 获取所有键
std::vector<std::string> MDictBase::keys() const {
    std::vector<std::string> result;
    for (const auto& key_pair : m_key_list) {
        result.push_back(key_pair.second);
    }
    return result;
}

Record MDictBase::get_offset_and_size(const string& word) {
    uint64_t offset = 0;
    uint32_t size = 0;

    KeyBlockIndex found;
    uint64_t pos = 0;
    for (auto p : m_key_block_indexs) {
        if (word.compare(p.first_key) >= 0 && word.compare(p.last_key) <= 0) {
            // 找到了key所在的block
            found = p;
            break;
        }
        pos += p.compressed_size;
        size++;
    }
    if (size == m_key_block_indexs.size()) return {0, 0};

    // 二分查找序号i所对应的key_block
    vector<uint8_t> key_block;
    std::ifstream file(m_filename, std::ios::binary);
    file.seekg(m_key_block_offset + pos);
    key_block = BinaryIO::read_bytes(file, found.compressed_size);
    file.close();
    auto keys = _split_key_block(key_block);
    for (int i = 0; i < keys.size(); i++) {
        if (keys[i].key == word) {
            offset = keys[i].offset_for_record_block;
            size = keys[i + 1].offset_for_record_block - offset;
        }
    }

    return {offset, size};
}

// 解析头部文本
std::map<std::string, std::string> MDictBase::_parse_header(
    const std::string& header_text) {
    std::map<std::string, std::string> tags;
    // std::istringstream stream(header_text);
    std::string line;

    // 使用正则表达式匹配 key="value" 形式
    std::regex pattern("(\\w+)=\"(.*?)\"");
    std::smatch matches;

    // 用于存储键值对的哈希表

    // 在字符串中查找所有匹配的键值对
    std::string::const_iterator searchStart(header_text.cbegin());
    while (
        std::regex_search(searchStart, header_text.cend(), matches, pattern)) {
        // matches[1] 是键，matches[2] 是值
        tags[matches[1]] = matches[2];
        // cout << "=========>" << matches[1] << matches[2];

        // 更新查找的位置
        searchStart = matches.suffix().first;
    }

    return tags;
}

// 读取文件头
std::map<std::string, std::string> MDictBase::_read_header() {
    std::ifstream file(m_filename, std::ios::binary);
    if (!file) {
        throw std::runtime_error("Failed to open file: " + m_filename);
    }

    // 读取头部大小
    uint32_t header_bytes_size = BinaryIO::read_value<uint32_t>(file);
    // header_bytes_size = BinaryIO::byte_swap(header_bytes_size, true);
    // std::cout << header_bytes_size << std::endl;

    // 读取头部数据
    std::vector<uint8_t> header_bytes =
        BinaryIO::read_bytes(file, header_bytes_size);

    // 读取并验证校验和
    uint32_t adler32 = BinaryIO::read_value<uint32_t>(file, false);
    // adler32 = BinaryIO::byte_swap(adler32, false);
    uint32_t calculated_adler = CompressionUtils::adler32(header_bytes);
    if (adler32 != calculated_adler) {
        throw std::runtime_error("Header checksum mismatch");
    }

    // 键块偏移
    m_key_block_offset = file.tellg();
    file.close();

    // 解析头部文本
    std::string header_text;
    // 处理UTF-16编码（假设是小端序）
    if (header_bytes.size() >= 2 &&
        header_bytes[header_bytes.size() - 2] == 0x00 &&
        header_bytes[header_bytes.size() - 1] == 0x00) {
        // UTF-16编码，需要转换为UTF-8
        // 简化实现：直接跳过零字节
        for (size_t i = 0; i < header_bytes.size(); i += 2) {
            if (header_bytes[i] != 0) {
                header_text.push_back(header_bytes[i]);
            }
        }
    } else {
        // 假设是UTF-8或其他单字节编码
        header_text.assign(header_bytes.begin(), header_bytes.end() - 1);
    }

    // 解析头部标签
    // cout << header_text << endl;
    std::map<std::string, std::string> header_tags = _parse_header(header_text);

    // 提取版本信息
    if (header_tags.find("GeneratedByEngineVersion") != header_tags.end()) {
        m_version = std::stof(header_tags["GeneratedByEngineVersion"]);
    } else {
        m_version = 1.0f;  // 默认版本
    }
    // std::cout << "字典生成器的版本: " << m_version << std::endl;

    // 设置数字宽度和格式
    if (m_version < 2.0) {
        m_number_width = 4;
        m_number_format = ">I";
    } else {
        m_number_width = 8;
        m_number_format = ">Q";
        // v3版本强制使用UTF-8
        if (m_version >= 3) {
            m_encoding = "UTF-8";
        }
    }

    // 处理加密标志
    if (header_tags.find("Encrypted") == header_tags.end() ||
        header_tags["Encrypted"] == "No") {
        m_encrypt = 0;
    } else if (header_tags["Encrypted"] == "Yes") {
        m_encrypt = 1;
    } else {
        m_encrypt = std::stoi(header_tags["Encrypted"]);
    }

    // 处理样式表
    if (header_tags.find("StyleSheet") != header_tags.end()) {
        // 简化实现，实际应用中需要更复杂的样式表解析
        m_stylesheet["default"] = {"text/css", header_tags["StyleSheet"]};
    }

    if (header_tags.find("Encoding") != header_tags.end()) {
        m_encoding = header_tags["Encoding"];
    }

    return header_tags;
}

// 读取keys
std::vector<std::pair<uint64_t, std::string>> MDictBase::_read_keys() {
    if (m_version >= 3) {
        return _read_keys_v3();
    } else {
        // 如果有加密且没有提供密码，尝试暴力破解（仅适用于engine <= 2）
        if ((m_encrypt & 0x01) && m_encrypted_key == nullptr) {
            std::cout << "Try Brutal Force on Encrypted Key Blocks"
                      << std::endl;
            return _read_keys_brutal();
        } else {
            return _read_keys_v1v2();
        }
    }
}

// v1/v2版本keys读取

// 键块的存储结构:
// key block header(40B/16B)
// key block info
// key block 1
// key block 2
// key block 3
// ...
std::vector<std::pair<uint64_t, std::string>> MDictBase::_read_keys_v1v2() {
    std::vector<std::pair<uint64_t, std::string>> key_list;

    std::ifstream file(m_filename, std::ios::binary);
    if (!file) {
        throw std::runtime_error("Failed to open file: " + m_filename);
    }

    file.seekg(m_key_block_offset);
    auto pos = file.tellg();

    // 1. 键块头部区域
    // version>=2为5x8(int64)个字节,
    // version<2为4x4(int32)个字节
    //
    std::vector<uint8_t> key_block_header;
    char header[5 * 8];
    if (m_version >= 2.0) {
        file.read(header, 5 * 8);
        key_block_header.assign((uint8_t*)header, (uint8_t*)header + 40);
    } else {
        file.read(header, 4 * 4);
        key_block_header.assign((uint8_t*)header, (uint8_t*)header + 16);
    }
    file.seekg(pos);

    int num_key_blocks = (m_version >= 2.0)
                             ? BinaryIO::read_value<uint64_t>(file)
                             : BinaryIO::read_value<uint32_t>(file);
    int num_entries = (m_version >= 2.0) ? BinaryIO::read_value<uint64_t>(file)
                                         : BinaryIO::read_value<uint32_t>(file);
    if (m_version >= 2.0) {
        int key_block_info_decomp_size = BinaryIO::read_value<uint64_t>(file);
    }
    int key_block_info_size = (m_version >= 2.0)
                                  ? BinaryIO::read_value<uint64_t>(file)
                                  : BinaryIO::read_value<uint32_t>(file);
    int key_block_size = (m_version >= 2.0)
                             ? BinaryIO::read_value<uint64_t>(file)
                             : BinaryIO::read_value<uint32_t>(file);

    m_num_entries = num_entries;
    cout << "key blocks的区块数量: " << num_key_blocks << endl;
    cout << "key的数量:" << num_entries << endl;
    cout << "key block info区域的字节数: " << key_block_info_size << endl;
    cout << "key blocks区域的字节数: " << key_block_size << endl;
    if (m_version >= 2.0) {
        int adler32_checksum = BinaryIO::read_value<uint32_t>(file);
        uint32_t calculated_adler = CompressionUtils::adler32(key_block_header);
        // cout << adler32_checksum << "\n" << calculated_adler << endl;
        cout << "key block header校验和验证(adler32): "
             << (adler32_checksum == calculated_adler) << endl;
    }
    // 2. 读取+解析key block info区域
    auto key_block_info = BinaryIO::read_bytes(file, key_block_info_size);
    vector<pair<uint64_t, uint64_t>> info_list =
        _decode_key_block_info(key_block_info);
    cout << "info list size: " << info_list.size() << endl;
    if (info_list.size() != num_key_blocks) {
        throw std::runtime_error("key block info size mismatch");
    }

    // 3. 读取+解析key block区域
    std::vector<uint8_t> key_blocks =
        BinaryIO::read_bytes(file, key_block_size);

    _decode_key_block(key_blocks, info_list);

    m_record_block_offset = file.tellg();
    file.close();
    return key_list;
}

// v3版本密钥读取
std::vector<std::pair<uint64_t, std::string>> MDictBase::_read_keys_v3() {
    std::vector<std::pair<uint64_t, std::string>> key_list;

    // 打开文件
    std::ifstream file(m_filename, std::ios::binary);
    if (!file) {
        throw std::runtime_error("Failed to open file: " + m_filename);
    }

    // 定位到键块偏移量
    file.seekg(m_key_block_offset);

    // 查找所有块的偏移量
    while (true) {
        uint32_t block_type = BinaryIO::read_value<uint32_t>(file);
        uint64_t block_size = BinaryIO::read_value<uint64_t>(file);
        long block_offset = file.tellg();

        // 根据块类型记录偏移量
        switch (block_type) {
            case 0x01000000:  // 记录数据
                m_record_block_offset = block_offset;
                break;
            case 0x02000000:  // 记录索引
                m_record_index_offset = block_offset;
                break;
            case 0x03000000:  // 键数据
                m_key_data_offset = block_offset;
                break;
            case 0x04000000:  // 键索引
                m_key_index_offset = block_offset;
                break;
            default:
                throw std::runtime_error("Unknown block type: " +
                                         std::to_string(block_type));
        }

        // 跳到下一个块
        file.seekg(block_size, std::ios::cur);

        // 检查是否到达文件末尾
        char test[4];
        if (file.read(test, 4).gcount() < 4) {
            break;
        }
        file.seekg(-4, std::ios::cur);
    }

    // 读取键数据
    file.seekg(m_key_data_offset);
    uint32_t number = BinaryIO::read_value<uint32_t>(file);
    uint64_t total_size = BinaryIO::read_value<uint64_t>(file);

    // 处理每个键块
    for (uint32_t i = 0; i < number; ++i) {
        uint32_t decompressed_size = BinaryIO::read_value<uint32_t>(file);
        uint32_t compressed_size = BinaryIO::read_value<uint32_t>(file);

        std::vector<uint8_t> block_data =
            BinaryIO::read_bytes(file, compressed_size);
        std::vector<uint8_t> decompressed_block_data =
            _decode_block(block_data, decompressed_size);

        // 解析键块
        auto keys_in_block = _split_key_block(decompressed_block_data);
        // key_list.insert(key_list.end(), keys_in_block.begin(),
        //                 keys_in_block.end());
    }

    file.close();
    m_num_entries = key_list.size();
    return key_list;
}

// 暴力破解密钥（简化版）
std::vector<std::pair<uint64_t, std::string>> MDictBase::_read_keys_brutal() {
    // 简化实现，实际应用中需要更复杂的暴力破解逻辑
    // 这里仅作为示例框架
    std::cout << "Brutal force is not implemented in this simplified version."
              << std::endl;
    return _read_keys_v1v2();
}

// 解码数据块
std::vector<uint8_t> MDictBase::_decode_block(const std::vector<uint8_t>& block,
                                              uint32_t decompressed_size) {
    if (block.empty()) {
        return {};
    }

    // 解析块信息[0:4], 小端
    uint32_t info = *reinterpret_cast<const uint32_t*>(block.data());
    info = BinaryIO::byte_swap(info, false);
    uint8_t compression_method = info & 0xf;
    uint8_t encryption_method = (info >> 4) & 0xf;
    uint8_t encryption_size = (info >> 8) & 0xff;

    // 解析Adler32校验和[4:8]
    uint32_t adler32 = *reinterpret_cast<const uint32_t*>(block.data() + 4);
    adler32 = BinaryIO::byte_swap(adler32);

    // 获取加密密钥
    std::vector<uint8_t> encrypted_key;
    if (m_encrypted_key) {
        encrypted_key.assign(m_encrypted_key,
                             m_encrypted_key + 32);  // 假设密钥长度为32字节
    } else {
        // 使用Adler32生成密钥
        std::vector<uint8_t> adler_bytes(
            reinterpret_cast<const uint8_t*>(&adler32),
            reinterpret_cast<const uint8_t*>(&adler32) + 4);
        encrypted_key = CryptoUtils::ripemd128(adler_bytes);
    }

    // 解密数据[8:]
    const vector<uint8_t> data(block.begin() + 8, block.end());

    std::vector<uint8_t> decrypted_block;
    if (encryption_method == 0) {
        decrypted_block = data;
    } else if (encryption_method == 1) {
        // Fast解密
        std::vector<uint8_t> encrypted_part(data.begin(),
                                            data.begin() + encryption_size);
        std::vector<uint8_t> decrypted_part =
            CryptoUtils::fast_decrypt(encrypted_part, encrypted_key);
        decrypted_block.insert(decrypted_block.end(), decrypted_part.begin(),
                               decrypted_part.end());
        decrypted_block.insert(decrypted_block.end(),
                               data.begin() + encryption_size, data.end());
    } else if (encryption_method == 2) {
        // Salsa20解密
        std::vector<uint8_t> encrypted_part(data.begin(),
                                            data.begin() + encryption_size);
        std::vector<uint8_t> decrypted_part =
            CryptoUtils::salsa20_decrypt(encrypted_part, encrypted_key);
        decrypted_block.insert(decrypted_block.end(), decrypted_part.begin(),
                               decrypted_part.end());
        decrypted_block.insert(decrypted_block.end(),
                               data.begin() + encryption_size, data.end());
    } else {
        throw std::runtime_error("Unsupported encryption method: " +
                                 std::to_string(encryption_method));
    }

    // 验证解密后的校验和（v3版本）
    if (m_version >= 3) {
        uint32_t calculated_adler = CompressionUtils::adler32(decrypted_block);
        if (adler32 != calculated_adler) {
            throw std::runtime_error("Decrypted block checksum mismatch");
        }
    }

    // 解压缩数据
    std::vector<uint8_t> decompressed_block;

    if (compression_method == 0) {
        decompressed_block = decrypted_block;
    } else if (compression_method == 1) {
        // LZO解压缩
        decompressed_block = CompressionUtils::lzo_decompress(
            decrypted_block, decompressed_size);
    } else if (compression_method == 2) {
        // ZLib解压缩
        decompressed_block = CompressionUtils::zlib_decompress(decrypted_block);
    } else {
        throw std::runtime_error("Unsupported compression method: " +
                                 std::to_string(compression_method));
    }

    // 验证解压后的校验和（v1/v2版本）
    if (m_version < 3) {
        uint32_t calculated_adler =
            CompressionUtils::adler32(decompressed_block);
        if (adler32 != calculated_adler) {
            throw std::runtime_error("Decompressed block checksum mismatch");
        }
    }

    return decompressed_block;
}

vector<pair<uint64_t, uint64_t>> MDictBase::_decode_key_block_info(
    vector<uint8_t> key_block_info_compressed) {
    // 校验block info头部数据
    vector<uint8_t> key_block_info;
    if (m_version >= 2) {
        // zlib压缩过了
        array<uint8_t, 4> data{0x02, 0x00, 0x00, 0x00};
        assert(!std::memcmp(key_block_info_compressed.data(), data.data(), 4));
        // 解密
        const auto& first = key_block_info_compressed.begin();
        const auto& end = key_block_info_compressed.end();
        vector<uint8_t> decrypted;
        // 如果加密了，则解密
        if (m_encrypt & 0x02) {
            vector<uint8_t> bytes(first + 4, first + 8);
            uint32_t num = 0x3695;
            num = BinaryIO::byte_swap(num, false);
            uint8_t* p = reinterpret_cast<uint8_t*>(&num);
            bytes.insert(bytes.end(), p, p + 4);
            auto key = CryptoUtils::ripemd128(bytes);
            decrypted.insert(decrypted.end(), first, first + 8);
            auto dec =
                CryptoUtils::fast_decrypt(vector<uint8_t>(first + 8, end), key);
            decrypted.insert(decrypted.end(), dec.begin(), dec.end());
        } else {
            decrypted = key_block_info_compressed;
        }
        // 解压缩
        vector<uint8_t> zin(decrypted.begin() + 8, decrypted.end());
        key_block_info = CompressionUtils::zlib_decompress(zin);
        // adler32 校验
        uint32_t adler32 =
            BinaryIO::unpack<uint32_t>(key_block_info_compressed, 4);
        cout << "adler32: " << adler32 << endl;
        uint32_t calculated_adler = CompressionUtils::adler32(key_block_info);
        if (adler32 != calculated_adler) {
            throw std::runtime_error(
                "Decompressed key block info checksum mismatch");
        }

    } else {
        key_block_info = key_block_info_compressed;
    }

    cout << "begin decode key block info" << endl;
    // decode
    vector<pair<uint64_t, uint64_t>> info_list;
    int num_entries = 0, i = 0, byte_width, text_term;
    if (m_version >= 2) {
        byte_width = 2;
        text_term = 1;
    } else {
        byte_width = 1;
        text_term = 0;
    }

    while (i < key_block_info.size()) {
        if (m_number_width == 8) {
            num_entries += BinaryIO::unpack<uint64_t>(key_block_info, i, true);
        } else if (m_number_width == 4) {
            num_entries += BinaryIO::unpack<uint32_t>(key_block_info, i, true);
        }
        i += m_number_width;
        // text head size
        int text_head_size =
            m_version >= 2 ? BinaryIO::unpack<uint16_t>(key_block_info, i, true)
                           : BinaryIO::unpack<uint8_t>(key_block_info, i, true);
        i += byte_width;
        int j = i;
        // text head
        m_encoding != "UTF-16" ? i += text_head_size + text_term
                               : i += (text_head_size + text_term) * 2;

        string first_key(key_block_info.cbegin() + j,
                         key_block_info.cbegin() + i);

        // text tail size
        int text_tail_size = m_version >= 2
                                 ? BinaryIO::unpack<uint16_t>(key_block_info, i)
                                 : BinaryIO::unpack<uint8_t>(key_block_info, i);
        i += byte_width;
        j = i;
        // text tail
        m_encoding != "UTF-16" ? i += text_tail_size + text_term
                               : i += (text_tail_size + text_term) * 2;
        string last_key(key_block_info.cbegin() + j,
                        key_block_info.cbegin() + i);
        // key block compressed size
        uint64_t compressed_size =
            m_number_width == 8 ? BinaryIO::unpack<uint64_t>(key_block_info, i)
                                : BinaryIO::unpack<uint32_t>(key_block_info, i);
        i += m_number_width;

        uint64_t decompressed_size =
            m_number_width == 8 ? BinaryIO::unpack<uint64_t>(key_block_info, i)
                                : BinaryIO::unpack<uint32_t>(key_block_info, i);
        i += m_number_width;
        // cout << compressed_size << " ," << decompressed_size << endl;
        info_list.emplace_back(compressed_size, decompressed_size);
        m_key_block_indexs.emplace_back(KeyBlockIndex{
            first_key, last_key, compressed_size, decompressed_size});
        // cout << first_key << " , " << last_key << endl;
    }
    return info_list;
}

// 解析key blocks, 一个key blocks包含多个key block
// @param key_block_compressed 键块儿数据
// @param info_list 键块儿信息列表[(键块压缩大小,键块解压后大小), ...]
// @return 所有键组成的列表[(键在记录块中的全局偏移量, 键),...]
vector<pair<uint64_t, string>> MDictBase::_decode_key_block(
    vector<uint8_t> key_block_compressed,
    vector<pair<uint64_t, uint64_t>> info_list) {
    vector<pair<uint64_t, string>> key_list;
    int i = 0;

    for (auto& it : info_list) {
        uint64_t& compressed_size = it.first;
        uint64_t& decompressed_size = it.second;

        // 切分键块
        const auto& it_head = key_block_compressed.cbegin();
        // block.insert(block.end(), it_head + i, it_head + i +
        // compressed_size);
        vector<uint8_t> block(it_head + i, it_head + i + compressed_size);

        // 解析key block,获取有效数据字段
        auto key_block = _decode_block(block, decompressed_size);
        // 提取一个key block中包含的所有key到一个列表中
        _split_key_block(key_block);
        i += compressed_size;
        // key_list.insert(key_list.end(), list.begin(), list.end());
    }

    return key_list;
}

// 解析key block，一个key block中包含多个key
// pair<int, string>中，第一个是key在record block中的偏移量
vector<KeyIndex> MDictBase::_split_key_block(const vector<uint8_t>& block) {
    std::vector<KeyIndex> key_list;

    size_t pos = 0;
    vector<uint8_t> end_chars;
    if (m_encoding == "UTF-16") {
        end_chars.assign({0x00, 0x00});
    } else {
        end_chars.push_back(0x00);
    }
    int width = end_chars.size();
    int num = 0;

    while (pos < block.size()) {
        // 读取key对应的记录在记录块中的偏移量
        uint64_t offset = 0;
        if (m_version >= 2.0) {
            if (pos + 8 > block.size()) break;
            // offset = *reinterpret_cast<const uint64_t*>(block.data() + pos);
            offset = BinaryIO::unpack<uint64_t>(block, pos);
        } else {
            if (pos + 4 > block.size()) break;
            // offset = *reinterpret_cast<const uint32_t*>(block.data() + pos);
            offset = BinaryIO::unpack<uint32_t>(block, pos);
        }

        int i = pos + m_number_width;
        int end_indx = 0;
        while (i < block.size()) {
            vector<uint8_t> tmp(block.begin() + i, block.begin() + i + width);
            if (!std::memcmp(tmp.data(), end_chars.data(), width)) {
                end_indx = i;
                break;
            }
            i += width;
        }
        string key(block.cbegin() + pos + m_number_width,
                   block.cbegin() + end_indx);

        pos = end_indx + width;

        // 添加到结果列表
        key_list.push_back(KeyIndex{key, offset, num});
        num++;
    }

    return key_list;
}
